{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "41aMjXA5rKoL"
      },
      "outputs": [],
      "source": [
        "# Copyright 2021 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bcKMGrqDrKoS"
      },
      "source": [
        "# Feedback or issues?\n",
        "\n",
        "For any feedback or questions, please open an [issue](https://github.com/googleapis/python-aiplatform/issues)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ij53c3EzrKoS"
      },
      "source": [
        "# Vertex SDK for Python: Custom Training using Python Package, Managed Text Dataset, and TF-Serving Container Example\n",
        "To use this Jupyter notebook, copy the notebook to a Google Cloud Notebooks instance with Tensorflow installed and open it. You can run each step, or cell, and see its results. To run a cell, use Shift+Enter. Jupyter automatically displays the return value of the last line in each cell. For more information about running notebooks in Google Cloud Notebook, see the [Google Cloud Notebook guide](https://cloud.google.com/vertex-ai/docs/general/notebooks).\n",
        "\n",
        "This notebook demonstrate how to create a Custom Model using Custom Python Package Training, with a Vertex AI Dataset, and how to serve the model using Tensorflow-Serving Container for online prediction, and batch prediction. It will require you provide a bucket where the dataset will be stored.\n",
        "\n",
        "Note: you may incur charges for training, prediction, storage or usage of other GCP products in connection with testing this SDK."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xOMNWzTbftDr"
      },
      "source": [
        "### Install Vertex SDK for Python\n",
        "\n",
        "\n",
        "After the SDK installation the kernel will be automatically restarted."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Be020jY-ftDv"
      },
      "outputs": [],
      "source": [
        "!pip3 uninstall -y google-cloud-aiplatform\n",
        "!pip3 install google-cloud-aiplatform\n",
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l6N5ru3xrKoU"
      },
      "source": [
        "### Enter Your Project and GCS Bucket\n",
        "\n",
        "Enter your Project Id in the cell below. Then run the cell to make sure the Cloud SDK uses the right project for all the commands in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d61oYG3KftDw"
      },
      "outputs": [],
      "source": [
        "MY_PROJECT = \"YOUR PROJECT ID\"\n",
        "MY_STAGING_BUCKET = \"gs://YOUR BUCKET\"  # bucket should be in same region as ucaip"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YHNj4P3vrKoV"
      },
      "source": [
        "### Set Your Application Name, Task Name, and Directories.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RnNL5St8rKoV"
      },
      "outputs": [],
      "source": [
        "APP_NAME = \"keras-text-class-stack-overflow-tag\"\n",
        "TASK_TYPE = \"mbsdk_custom-py-pkg-training\"\n",
        "\n",
        "TASK_NAME = f\"{TASK_TYPE}_{APP_NAME}\"\n",
        "\n",
        "TASK_DIR = f\"./{TASK_NAME}\"\n",
        "DATA_DIR = f\"{TASK_DIR}/data\"\n",
        "\n",
        "print(f\"Task Name:      {TASK_NAME}\")\n",
        "print(f\"Task Directory: {TASK_DIR}\")\n",
        "print(f\"Data Directory: {DATA_DIR}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9hZmDcfLS13W"
      },
      "source": [
        "### Set a GCS Prefix\n",
        "\n",
        "If you want to centeralize all input and output files under the gcs location."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W_4gzaP0SyQd"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = MY_STAGING_BUCKET.split(\"gs://\")[1]\n",
        "GCS_PREFIX = f\"{TASK_TYPE}/{APP_NAME}\"\n",
        "\n",
        "print(f\"Bucket Name:    {BUCKET_NAME}\")\n",
        "print(f\"GCS Prefix:     {GCS_PREFIX}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5T1d5uBoftDw"
      },
      "source": [
        "# Stack Overflow Data\n",
        "We download the stack overflow data from from  https://storage.googleapis.com/download.tensorflow.org/data/stack_overflow_16k.tar.gz and will create a Vertex AI managed text dataset. \n",
        "\n",
        "The Stack Overflow Data is licensed under the Creative Commons Attribution-ShareAlike 3.0 Unported License. To view a copy of this license, visit http://creativecommons.org/licenses/by-sa/3.0/ \n",
        "\n",
        "For more information about this dataset please visit: https://console.cloud.google.com/marketplace/details/stack-exchange/stack-overflow\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DJF047yNftDw"
      },
      "source": [
        "### Utility Functions to Download Data and Prepare CSV Files for Creating Vertex AI Managed Dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9yOl-l_oftDx"
      },
      "outputs": [],
      "source": [
        "import csv\n",
        "import os\n",
        "\n",
        "from google.cloud import storage\n",
        "from tensorflow.keras import utils\n",
        "\n",
        "\n",
        "def upload_blob(bucket_name, source_file_name, destination_blob_name):\n",
        "    \"\"\"Uploads a file to the bucket.\"\"\"\n",
        "\n",
        "    storage_client = storage.Client()\n",
        "    bucket = storage_client.bucket(bucket_name)\n",
        "    blob = bucket.blob(destination_blob_name)\n",
        "\n",
        "    blob.upload_from_filename(source_file_name)\n",
        "\n",
        "    destination_file_name = os.path.join(\"gs://\", bucket_name, destination_blob_name)\n",
        "\n",
        "    return destination_file_name\n",
        "\n",
        "\n",
        "def download_data(data_dir):\n",
        "    \"\"\"Download data.\"\"\"\n",
        "\n",
        "    if not os.path.exists(data_dir):\n",
        "        os.makedirs(data_dir)\n",
        "\n",
        "    url = \"https://storage.googleapis.com/download.tensorflow.org/data/stack_overflow_16k.tar.gz\"\n",
        "    dataset = utils.get_file(\n",
        "        \"stack_overflow_16k.tar.gz\",\n",
        "        url,\n",
        "        untar=True,\n",
        "        cache_dir=data_dir,\n",
        "        cache_subdir=\"\",\n",
        "    )\n",
        "    data_dir = os.path.join(os.path.dirname(dataset))\n",
        "\n",
        "    return data_dir\n",
        "\n",
        "\n",
        "def upload_train_data_to_gcs(train_data_dir, bucket_name, destination_blob_prefix):\n",
        "    \"\"\"Create CSV file using train data content.\"\"\"\n",
        "\n",
        "    train_data_dir = os.path.join(data_dir, \"train\")\n",
        "    train_data_fn = os.path.join(data_dir, \"train.csv\")\n",
        "\n",
        "    fp = open(train_data_fn, \"w\", encoding=\"utf8\")\n",
        "    writer = csv.writer(\n",
        "        fp, delimiter=\",\", quotechar='\"', quoting=csv.QUOTE_ALL, lineterminator=\"\\n\"\n",
        "    )\n",
        "\n",
        "    for root, _, files in os.walk(train_data_dir):\n",
        "        for file in files:\n",
        "            if file.endswith(\".txt\"):\n",
        "                class_name = root.split(\"/\")[-1]\n",
        "                file_fn = os.path.join(root, file)\n",
        "                with open(file_fn, \"r\") as f:\n",
        "                    content = f.readlines()\n",
        "                    lines = [x.strip().strip('\"') for x in content]\n",
        "                    writer.writerow((lines[0], class_name))\n",
        "\n",
        "    fp.close()\n",
        "\n",
        "    train_gcs_url = upload_blob(\n",
        "        bucket_name, train_data_fn, os.path.join(destination_blob_prefix, \"train.csv\")\n",
        "    )\n",
        "\n",
        "    return train_gcs_url"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q3UDGt8MrKoY"
      },
      "source": [
        "### Download Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ks2ahWfUrKoa"
      },
      "outputs": [],
      "source": [
        "data_dir = download_data(DATA_DIR)\n",
        "print(f\"Data is downloaded to: {data_dir}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l7AHtjj3rKoa"
      },
      "outputs": [],
      "source": [
        "!ls $data_dir"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "puroD_Wr-W4Y"
      },
      "outputs": [],
      "source": [
        "!ls $data_dir/train"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bQNL5AjtrKoa"
      },
      "source": [
        "### Prepare CSV Files for Creating Managed Dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9ewrG6VArKob"
      },
      "source": [
        "#### Create CSV Files using Data Content"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GvbpGr0ArKob"
      },
      "outputs": [],
      "source": [
        "gcs_source_train_url = upload_train_data_to_gcs(\n",
        "    train_data_dir=os.path.join(data_dir, \"train\"),\n",
        "    bucket_name=BUCKET_NAME,\n",
        "    destination_blob_prefix=f\"{GCS_PREFIX}/data\",\n",
        ")\n",
        "\n",
        "print(f\"Train data content is loaded to {gcs_source_train_url}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Gg_OiKJUrKoc"
      },
      "outputs": [],
      "source": [
        "!gsutil ls gs://$BUCKET_NAME/$GCS_PREFIX/data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qzKx7Z2WrKoc"
      },
      "source": [
        "# Create Custom Training Python Package\n",
        "\n",
        "Before you can perform custom training with a pre-built container, you must create a [Python Source Distribution](https://docs.python.org/3/distutils/sourcedist.html) that contains your training application and upload it to a Cloud Storage bucket that your Google Cloud project can access.\n",
        "\n",
        "We will create a directory and write all of our package build artifacts into that folder."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eBDjNpkLrKoc"
      },
      "outputs": [],
      "source": [
        "PYTHON_PACKAGE_APPLICATION_DIR = f\"{TASK_NAME}/trainer\"\n",
        "\n",
        "!mkdir -p $PYTHON_PACKAGE_APPLICATION_DIR\n",
        "!touch $PYTHON_PACKAGE_APPLICATION_DIR/__init__.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sPSRvSvMrKod"
      },
      "source": [
        "### Write the Training Script"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zupq7mwzrKod"
      },
      "outputs": [],
      "source": [
        "%%writefile {PYTHON_PACKAGE_APPLICATION_DIR}/task.py\n",
        "\n",
        "\n",
        "import os\n",
        "import argparse\n",
        "\n",
        "import tensorflow as tf\n",
        "from tensorflow.keras import layers\n",
        "from tensorflow.keras import losses\n",
        "from tensorflow.keras.layers.experimental.preprocessing import TextVectorization\n",
        "\n",
        "import json\n",
        "import tqdm\n",
        "\n",
        "VOCAB_SIZE = 10000\n",
        "MAX_SEQUENCE_LENGTH = 250\n",
        "\n",
        "def str2bool(v):\n",
        "  if isinstance(v, bool):\n",
        "    return v\n",
        "  if v.lower() in ('yes', 'true', 't', 'y', '1'):\n",
        "    return True\n",
        "  elif v.lower() in ('no', 'false', 'f', 'n', '0'):\n",
        "    return False\n",
        "  else:\n",
        "    raise argparse.ArgumentTypeError('Boolean value expected.')\n",
        "\n",
        "def build_model(num_classes, loss, optimizer, metrics, vectorize_layer):\n",
        "  # vocab_size is VOCAB_SIZE + 1 since 0 is used additionally for padding.\n",
        "  model = tf.keras.Sequential([\n",
        "      vectorize_layer,\n",
        "      layers.Embedding(VOCAB_SIZE + 1, 64, mask_zero=True),\n",
        "      layers.Conv1D(64, 5, padding=\"valid\", activation=\"relu\", strides=2),\n",
        "      layers.GlobalMaxPooling1D(),\n",
        "      layers.Dense(num_classes),\n",
        "      layers.Activation('softmax')\n",
        "  ])\n",
        "  model.compile(\n",
        "      loss=loss,\n",
        "      optimizer=optimizer,\n",
        "      metrics=metrics)\n",
        "\n",
        "  return model\n",
        "\n",
        "def get_string_labels(predicted_scores_batch, class_names):\n",
        "  predicted_labels = tf.argmax(predicted_scores_batch, axis=1)\n",
        "  predicted_labels = tf.gather(class_names, predicted_labels)\n",
        "  return predicted_labels\n",
        "\n",
        "def predict(export_model, class_names, inputs):\n",
        "  predicted_scores = export_model.predict(inputs)\n",
        "  predicted_labels = get_string_labels(predicted_scores, class_names)\n",
        "  return predicted_labels\n",
        "\n",
        "def parse_args():\n",
        "  parser = argparse.ArgumentParser(\n",
        "      description='Keras Text Classification on Stack Overflow Questions')\n",
        "  parser.add_argument(\n",
        "      '--epochs', default=25, type=int, help='number of training epochs')\n",
        "  parser.add_argument(\n",
        "      '--batch-size', default=16, type=int, help='mini-batch size')\n",
        "  parser.add_argument(\n",
        "      '--model-dir', default=os.getenv('AIP_MODEL_DIR'), type=str, help='model directory')\n",
        "  parser.add_argument(\n",
        "      '--data-dir', default='./data', type=str, help='data directory')\n",
        "  parser.add_argument(\n",
        "      '--test-run', default=False, type=str2bool, help='test run the training application, i.e. 1 epoch for training using sample dataset')\n",
        "  parser.add_argument(\n",
        "      '--model-version', default=1, type=int, help='model version')\n",
        "  args = parser.parse_args()\n",
        "  return args\n",
        "\n",
        "def load_aip_dataset(aip_data_uri_pattern, batch_size, class_names, test_run, shuffle=True, seed=42):\n",
        "\n",
        "  data_file_urls = list()\n",
        "  labels = list()\n",
        "\n",
        "  class_indices = dict(zip(class_names, range(len(class_names))))\n",
        "  num_classes = len(class_names)\n",
        "\n",
        "  for aip_data_uri in tqdm.tqdm(tf.io.gfile.glob(pattern=aip_data_uri_pattern)):\n",
        "    with tf.io.gfile.GFile(name=aip_data_uri, mode='r') as gfile:\n",
        "      for line in gfile.readlines():\n",
        "        line = json.loads(line)\n",
        "        data_file_urls.append(line['textContent'])\n",
        "        classification_annotation = line['classificationAnnotations'][0]\n",
        "        label = classification_annotation['displayName']\n",
        "        labels.append(class_indices[label])\n",
        "        if test_run:\n",
        "          break\n",
        "\n",
        "  data = list()\n",
        "  for data_file_url in tqdm.tqdm(data_file_urls):\n",
        "    with tf.io.gfile.GFile(name=data_file_url, mode='r') as gf:\n",
        "      txt = gf.read()\n",
        "      data.append(txt)\n",
        "\n",
        "  print(f' data files count: {len(data_file_urls)}')\n",
        "  print(f' data count: {len(data)}')\n",
        "  print(f' labels count: {len(labels)}')\n",
        "\n",
        "  dataset = tf.data.Dataset.from_tensor_slices(data)\n",
        "  label_ds = tf.data.Dataset.from_tensor_slices(labels)\n",
        "  label_ds = label_ds.map(lambda x: tf.one_hot(x, num_classes))\n",
        "\n",
        "  dataset = tf.data.Dataset.zip((dataset, label_ds))\n",
        "\n",
        "  if shuffle:\n",
        "    # Shuffle locally at each iteration\n",
        "    dataset = dataset.shuffle(buffer_size=batch_size * 8, seed=seed)\n",
        "  dataset = dataset.batch(batch_size)\n",
        "  # Users may need to reference `class_names`.\n",
        "  dataset.class_names = class_names\n",
        "\n",
        "  return dataset\n",
        "\n",
        "def main():\n",
        "\n",
        "  args = parse_args()\n",
        "\n",
        "  class_names = ['csharp', 'java', 'javascript', 'python']\n",
        "  class_indices = dict(zip(class_names, range(len(class_names))))\n",
        "  num_classes = len(class_names)\n",
        "  print(f' class names: {class_names}')\n",
        "  print(f' class indices: {class_indices}')\n",
        "  print(f' num classes: {num_classes}')\n",
        "\n",
        "  epochs = 1 if args.test_run else args.epochs\n",
        "\n",
        "  aip_model_dir = os.environ.get('AIP_MODEL_DIR')\n",
        "  aip_data_format = os.environ.get('AIP_DATA_FORMAT')\n",
        "  aip_training_data_uri = os.environ.get('AIP_TRAINING_DATA_URI')\n",
        "  aip_validation_data_uri = os.environ.get('AIP_VALIDATION_DATA_URI')\n",
        "  aip_test_data_uri = os.environ.get('AIP_TEST_DATA_URI')\n",
        "\n",
        "  print(f\"aip_model_dir: {aip_model_dir}\")\n",
        "  print(f\"aip_data_format: {aip_data_format}\")\n",
        "  print(f\"aip_training_data_uri: {aip_training_data_uri}\")\n",
        "  print(f\"aip_validation_data_uri: {aip_validation_data_uri}\")\n",
        "  print(f\"aip_test_data_uri: {aip_test_data_uri}\")\n",
        "\n",
        "  print('Loading AIP dataset')\n",
        "  train_ds = load_aip_dataset(\n",
        "      aip_training_data_uri, args.batch_size, class_names, args.test_run)\n",
        "  print('AIP training dataset is loaded')\n",
        "  val_ds = load_aip_dataset(\n",
        "      aip_validation_data_uri, 1, class_names, args.test_run)\n",
        "  print('AIP validation dataset is loaded')\n",
        "  test_ds = load_aip_dataset(\n",
        "      aip_test_data_uri, 1, class_names, args.test_run)\n",
        "  print('AIP test dataset is loaded')\n",
        "\n",
        "  vectorize_layer = TextVectorization(\n",
        "      max_tokens=VOCAB_SIZE,\n",
        "      output_mode='int',\n",
        "      output_sequence_length=MAX_SEQUENCE_LENGTH)\n",
        "\n",
        "  train_text = train_ds.map(lambda text, labels: text)\n",
        "  vectorize_layer.adapt(train_text)\n",
        "  print('The vectorize_layer is adapted')\n",
        "\n",
        "\n",
        "  print('Build model')\n",
        "  optimizer = 'adam'\n",
        "  metrics = ['accuracy']\n",
        "\n",
        "  model = build_model(\n",
        "      num_classes, losses.CategoricalCrossentropy(from_logits=True), optimizer, metrics, vectorize_layer)\n",
        "\n",
        "  history = model.fit(train_ds, validation_data=val_ds, epochs=epochs)\n",
        "  history = history.history\n",
        "\n",
        "  print('Training accuracy: {acc}, loss: {loss}'.format(\n",
        "      acc=history['accuracy'][-1], loss=history['loss'][-1]))\n",
        "  print('Validation accuracy: {acc}, loss: {loss}'.format(\n",
        "      acc=history['val_accuracy'][-1], loss=history['val_loss'][-1]))\n",
        "\n",
        "  loss, accuracy = model.evaluate(test_ds)\n",
        "  print('Test accuracy: {acc}, loss: {loss}'.format(\n",
        "      acc=accuracy, loss=loss))\n",
        "\n",
        "  inputs = [\n",
        "      \"how do I extract keys from a dict into a list?\",  # python\n",
        "      \"debug public static void main(string[] args) {...}\",  # java\n",
        "  ]\n",
        "  predicted_labels = predict(model, class_names, inputs)\n",
        "  for input, label in zip(inputs, predicted_labels):\n",
        "    print(f'Question: {input}')\n",
        "    print(f'Predicted label: {label.numpy()}')\n",
        "\n",
        "  model_export_path = os.path.join(args.model_dir, str(args.model_version))\n",
        "  model.save(model_export_path)\n",
        "  print(f'Model version {args.model_version} is exported to {args.model_dir}')\n",
        "\n",
        "  loaded = tf.saved_model.load(model_export_path)\n",
        "  input_name = list(loaded.signatures['serving_default'].structured_input_signature[1].keys())[0]\n",
        "  print(f'Serving function input: {input_name}')\n",
        "\n",
        "  return\n",
        "\n",
        "if __name__ == '__main__':\n",
        "  main()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAhynoNBrKoh"
      },
      "source": [
        "### Build Package"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kXSBQcqdrKof"
      },
      "outputs": [],
      "source": [
        "%%writefile {TASK_DIR}/setup.py\n",
        "\n",
        "from setuptools import find_packages\n",
        "from setuptools import setup\n",
        "\n",
        "setup(\n",
        "    name='trainer',\n",
        "    version='0.1',\n",
        "    packages=find_packages(),\n",
        "    install_requires=(),\n",
        "    include_package_data=True,\n",
        "    description='My training application.'\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MclanW0UrKoh"
      },
      "outputs": [],
      "source": [
        "!ls $TASK_DIR"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j19Nk2vzrKoh"
      },
      "outputs": [],
      "source": [
        "!cd $TASK_DIR && python3 setup.py sdist --formats=gztar"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9c1r_k69rKoi"
      },
      "outputs": [],
      "source": [
        "!ls -ltr $TASK_DIR/dist/trainer-0.1.tar.gz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zyQv9AcNrKoi"
      },
      "source": [
        "### Upload the Package to GCS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9WjF8SFrrKoi"
      },
      "outputs": [],
      "source": [
        "destination_blob_name = f\"custom-training-python-package/{APP_NAME}/trainer-0.1.tar.gz\"\n",
        "source_file_name = f\"{TASK_DIR}/dist/trainer-0.1.tar.gz\"\n",
        "\n",
        "python_package_gcs_uri = upload_blob(\n",
        "    BUCKET_NAME, source_file_name, destination_blob_name\n",
        ")\n",
        "python_module_name = \"trainer.task\"\n",
        "\n",
        "print(f\"Custom Training Python Package is uploaded to: {python_package_gcs_uri}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hWJpv4ejrWi0"
      },
      "source": [
        "# Create TensorFlow Serving Container"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IMOad4Z4aOgv"
      },
      "source": [
        "Download the TensorFlow Serving Docker image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KC8e1yrHrVyA"
      },
      "outputs": [],
      "source": [
        "!docker pull tensorflow/serving:latest"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rtiNicSyaTk1"
      },
      "source": [
        "Create a tag for registering the image and register the image with Cloud Container Registry (gcr.io)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sUfgSddcsRvw"
      },
      "outputs": [],
      "source": [
        "TF_SERVING_CONTAINER_IMAGE_URI = f\"gcr.io/{MY_PROJECT}/tf-serving\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bw_A7ynTsjTF"
      },
      "outputs": [],
      "source": [
        "!docker tag tensorflow/serving $TF_SERVING_CONTAINER_IMAGE_URI\n",
        "!docker push $TF_SERVING_CONTAINER_IMAGE_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rc_MIJwlrKoj"
      },
      "source": [
        "# Run Custom Python Package Training with Managed Text Dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eiNL0HSVrKoj"
      },
      "source": [
        "## Initialize Vertex SDK for Python\n",
        "\n",
        "Initialize the *client* for Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J5MqvGtMrKoj"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=MY_PROJECT, staging_bucket=MY_STAGING_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KGT1uT-HrKoj"
      },
      "source": [
        "## Create a Dataset on Vertex AI\n",
        "We will now create a Vertex AI text dataset using the previously prepared csv files. Choose one of the options below. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lY566K2crKok"
      },
      "outputs": [],
      "source": [
        "dataset_display_name = f\"temp-{APP_NAME}-content\"\n",
        "gcs_source = gcs_source_train_url"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kyLoUsx9rKok"
      },
      "source": [
        "#### Option 1: Create a Dataset with CSV File"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h7QwC1ThrKok"
      },
      "outputs": [],
      "source": [
        "dataset = aiplatform.TextDataset.create(\n",
        "    display_name=dataset_display_name,\n",
        "    gcs_source=gcs_source,\n",
        "    import_schema_uri=aiplatform.schema.dataset.ioformat.text.single_label_classification,\n",
        "    sync=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xksBj3M0rKok"
      },
      "source": [
        "#### Option 2: Create a Dataset, then Import CSV File"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "93w8LGlx7uWy"
      },
      "source": [
        "```\n",
        "dataset = aiplatform.TextDataset.create(\n",
        "    display_name=dataset_display_name,\n",
        ")\n",
        "dataset.import_data(\n",
        "    gcs_source=gcs_source, \n",
        "    import_schema_uri=aiplatform.schema.dataset.ioformat.text.single_label_classification,\n",
        "    sync=False\n",
        ")\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ncP9QwvkrKok"
      },
      "source": [
        "#### Option 3: Retrieve a Dataset on Vertex AI\n",
        "If you have previously created a Dataset on Vertex AI, you can retrieve the dataset using the `dataset_name`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OmiQDoS_7ncz"
      },
      "source": [
        "```\n",
        "dataset_name = 'YOUR DATASET NAME'\n",
        "\n",
        "dataset = aiplatform.TextDataset(dataset_name)\n",
        "dataset.resource_name\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NzLvubIrrKon"
      },
      "source": [
        "## Launch a Training Job and Create a Model on Vertex AI\n",
        "\n",
        "We will now train a model with the python package we just built."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cyz86ixRYNtT"
      },
      "source": [
        "### Config a Training Job"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9zk8SRuK1Xvi"
      },
      "outputs": [],
      "source": [
        "MODEL_NAME = APP_NAME\n",
        "PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI = (\n",
        "    \"gcr.io/cloud-aiplatform/training/tf-cpu.2-3:latest\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rxc_yVxkbbbd"
      },
      "source": [
        "You will need to specify the python package that was built and uploaded to GCS, the module name of the python package, the pre-built training container image uri for training, and in this example, we are using TensorFlow serving container for prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uRGrFdxOftD1"
      },
      "outputs": [],
      "source": [
        "job = aiplatform.CustomPythonPackageTrainingJob(\n",
        "    display_name=f\"temp_{TASK_NAME}_tf-serving\",\n",
        "    python_package_gcs_uri=python_package_gcs_uri,\n",
        "    python_module_name=python_module_name,\n",
        "    container_uri=PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI,\n",
        "    model_serving_container_image_uri=TF_SERVING_CONTAINER_IMAGE_URI,\n",
        "    model_serving_container_command=[\"/usr/bin/tensorflow_model_server\"],\n",
        "    model_serving_container_args=[\n",
        "        f\"--model_name={MODEL_NAME}\",\n",
        "        \"--model_base_path=$(AIP_STORAGE_URI)\",\n",
        "        \"--rest_api_port=8080\",\n",
        "        \"--port=8500\",\n",
        "        \"--file_system_poll_wait_seconds=31540000\",\n",
        "    ],\n",
        "    model_serving_container_predict_route=f\"/v1/models/{MODEL_NAME}:predict\",\n",
        "    model_serving_container_health_route=f\"/v1/models/{MODEL_NAME}\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JPfZGlZeTN72"
      },
      "source": [
        "### Run the Training Job"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V6iY5wcerKon"
      },
      "outputs": [],
      "source": [
        "model = job.run(\n",
        "    dataset=dataset,\n",
        "    annotation_schema_uri=aiplatform.schema.dataset.annotation.text.classification,\n",
        "    args=[\"--epochs\", \"50\"],\n",
        "    replica_count=1,\n",
        "    model_display_name=f\"temp_{TASK_NAME}_tf-serving\",\n",
        "    sync=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zV1PjANLrKoo"
      },
      "outputs": [],
      "source": [
        "model.wait()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XO0-zdXxrKoo"
      },
      "source": [
        "# Deploy a Model and Create an Endpoint on Vertex AI\n",
        "\n",
        "Deploy your model, then wait until the model FINISHES deployment before proceeding to prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tEg2IDwPftD2"
      },
      "outputs": [],
      "source": [
        "endpoint = model.deploy(machine_type=\"n1-standard-4\", sync=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9AoIJdCHrKop"
      },
      "outputs": [],
      "source": [
        "endpoint.wait()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z87MXJEDrKoq"
      },
      "source": [
        "## Predict on the Endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vQXGCg8IrKoq"
      },
      "outputs": [],
      "source": [
        "class_names = [\"csharp\", \"java\", \"javascript\", \"python\"]\n",
        "\n",
        "class_ids = range(len(class_names))\n",
        "\n",
        "class_indices = dict(zip(class_names, class_ids))\n",
        "class_maps = dict(zip(class_ids, class_names))\n",
        "print(f\"Class Indices: {class_indices}\")\n",
        "print(f\"Class Maps:    {class_maps}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Di8RtRxipm63"
      },
      "outputs": [],
      "source": [
        "text_inputs = [\n",
        "    \"how do I extract keys from a dict into a list?\",  # python\n",
        "    \"debug public static void main(string[] args) {...}\",  # java\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dPHYGAlvlTAj"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "predictions = endpoint.predict(instances=[[text] for text in text_inputs])\n",
        "for text, predicted_scores in zip(text_inputs, predictions.predictions):\n",
        "    class_id = np.argmax(predicted_scores)\n",
        "    class_name = class_maps[class_id]\n",
        "    print(f\"Question: {text}\")\n",
        "    print(f\"Predicted Tag: {class_name}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gRkvzr2-LgA9"
      },
      "source": [
        "# Batch Prediction Job on the Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nlzRioF7UGFd"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "\n",
        "def upload_test_data_to_gcs(test_data_dir, test_gcs_url):\n",
        "    \"\"\"Create JSON file using test data content.\"\"\"\n",
        "\n",
        "    input_name = \"text_vectorization_input\"\n",
        "\n",
        "    with tf.io.gfile.GFile(test_gcs_url, \"w\") as gf:\n",
        "\n",
        "        for root, _, files in os.walk(test_data_dir):\n",
        "            for file in files:\n",
        "                if file.endswith(\".txt\"):\n",
        "                    file_fn = os.path.join(root, file)\n",
        "                    with open(file_fn, \"r\") as f:\n",
        "                        content = f.readlines()\n",
        "                        lines = [x.strip().strip('\"') for x in content]\n",
        "\n",
        "                        data = {input_name: [lines[0]]}\n",
        "                        gf.write(json.dumps(data))\n",
        "                        gf.write(\"\\n\")\n",
        "    return"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C_sbOm-5ud5C"
      },
      "outputs": [],
      "source": [
        "gcs_source_test_url = f\"gs://{BUCKET_NAME}/{GCS_PREFIX}/data/test.json\"\n",
        "upload_test_data_to_gcs(\n",
        "    test_data_dir=os.path.join(data_dir, \"test\"), test_gcs_url=gcs_source_test_url\n",
        ")\n",
        "\n",
        "print(f\"Test data content is loaded to {gcs_source_test_url}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QBMBk2WxLqBP"
      },
      "outputs": [],
      "source": [
        "!gsutil ls $gcs_source_test_url"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2JhoTiD5LuDW"
      },
      "outputs": [],
      "source": [
        "batch_predict_job = model.batch_predict(\n",
        "    job_display_name=f\"temp_{TASK_NAME}_tf-serving\",\n",
        "    gcs_source=gcs_source_test_url,\n",
        "    gcs_destination_prefix=f\"gs://{BUCKET_NAME}/{GCS_PREFIX}/batch_prediction\",\n",
        "    machine_type=\"n1-standard-4\",\n",
        "    sync=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "retsI3LLls_W"
      },
      "outputs": [],
      "source": [
        "batch_predict_job.wait()\n",
        "bp_iter_outputs = batch_predict_job.iter_outputs()\n",
        "\n",
        "prediction_errors_stats = list()\n",
        "prediction_results = list()\n",
        "for blob in bp_iter_outputs:\n",
        "    if blob.name.split(\"/\")[-1].startswith(\"prediction.errors_stats\"):\n",
        "        prediction_errors_stats.append(blob.name)\n",
        "    if blob.name.split(\"/\")[-1].startswith(\"prediction.results\"):\n",
        "        prediction_results.append(blob.name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KmfK3Tzzlv9C"
      },
      "outputs": [],
      "source": [
        "tags = list()\n",
        "for prediction_result in prediction_results:\n",
        "    gfile_name = f\"gs://{bp_iter_outputs.bucket.name}/{prediction_result}\"\n",
        "    with tf.io.gfile.GFile(name=gfile_name, mode=\"r\") as gfile:\n",
        "        for line in gfile.readlines():\n",
        "            line = json.loads(line)\n",
        "            text = line[\"instance\"][\"text_vectorization_input\"][0]\n",
        "            prediction = line[\"prediction\"]\n",
        "            class_id = np.argmax(prediction)\n",
        "            class_name = class_maps[class_id]\n",
        "            tags.append([text, class_name])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UcMQ-daLn_oh"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "\n",
        "tags_df = pd.DataFrame(tags, columns=[\"question\", \"tag\"])\n",
        "tags_df.head()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QUV_PHjtoCpQ"
      },
      "outputs": [],
      "source": [
        "tags_df[\"tag\"].value_counts()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "bcKMGrqDrKoS",
        "q3UDGt8MrKoY",
        "bQNL5AjtrKoa",
        "9ewrG6VArKob",
        "eiNL0HSVrKoj",
        "KGT1uT-HrKoj",
        "kyLoUsx9rKok",
        "xksBj3M0rKok",
        "ncP9QwvkrKok",
        "Cyz86ixRYNtT"
      ],
      "name": "AI_Platform_(Unified)_SDK_Custom_Training_Python_Package_Managed_Text_Dataset_Tensorflow_Serving_Container.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
